from pymilvus import connections, Collection, FieldSchema, CollectionSchema, DataType
import numpy as np
from typing import List, Dict
import json

class RAGModel:
    def __init__(self, collection_name: str):
        self.collection_name = collection_name
        self._init_milvus()
    
    def _init_milvus(self):
        connections.connect(host="localhost", port="19530")
        
        fields = [
            FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
            FieldSchema(name="type", dtype=DataType.VARCHAR, max_length=20),
            FieldSchema(name="content", dtype=DataType.VARCHAR, max_length=65535),
            FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=512),
            FieldSchema(name="page", dtype=DataType.INT64)
        ]
        
        schema = CollectionSchema(fields=fields)
        self.collection = Collection(name=self.collection_name, schema=schema)
        
        # 创建索引
        index_params = {
            "metric_type": "L2",
            "index_type": "IVF_FLAT",
            "params": {"nlist": 1024}
        }
        self.collection.create_index(field_name="embedding", index_params=index_params)
    
    def add_documents(self, documents: List[Dict], embeddings: List[np.ndarray]):
        entities = []
        for doc, emb in zip(documents, embeddings):
            # 确保向量数据是浮点数列表
            embedding = emb.tolist() if isinstance(emb, np.ndarray) else emb
            if not isinstance(embedding, list):
                continue
            
            entities.append({
                "type": doc["type"],
                "content": str(doc["content"]),  # 确保content是字符串
                "embedding": embedding,
                "page": int(doc["page"])  # 确保page是整数
            })
        
        if entities:
            self.collection.insert(entities)
            self.collection.flush()
    
    def search(self, query_embedding: np.ndarray, top_k: int = 5) -> List[Dict]:
        search_params = {"metric_type": "L2", "params": {"nprobe": 10}}
        results = self.collection.search(
            data=[query_embedding.tolist()],
            anns_field="embedding",
            param=search_params,
            limit=top_k,
            output_fields=["type", "content", "page"]
        )
        
        # 验证搜索结果
        valid_results = []
        for hit in results[0]:
            entity = hit.entity
            # 检查content是否为有效字符串
            if not isinstance(entity.get('content'), str) or not entity.get('content').strip():
                continue
            # 检查page是否为有效数字
            if not isinstance(entity.get('page'), (int, float)):
                continue
            valid_results.append(entity)
        
        return valid_results 